import config
from pathlib import Path
from obs import ObsClient
from pcdet.config import cfg, cfg_from_yaml_file
from pcdet.models import build_network
from ultralytics import YOLO
from annotate import DemoDataset, annotate, data_path_dict
import logging
import logging.handlers
import sys

logger = logging.getLogger("my_log")
logger.setLevel(logging.INFO)
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
file_handler = logging.handlers.RotatingFileHandler("my_log.log", maxBytes=10**6, backupCount=3)
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
console_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)
logger.addHandler(console_handler)
logger.addHandler(file_handler)


if __name__ == "__main__":
    try:
        #in input param, there is datasetId, datasetType, datasetPath, imageConfTh, pcdConfTh
        if len(sys.argv) != 7:
            error_message = "parameter number is {}, need 7 including data_annotation.py".format(len(sys.argv))
            logger.error(error_message)
            raise Exception(error_message)
        datasetId = sys.argv[1]
        datasetType = sys.argv[2]
        datasetPath = sys.argv[3]
        imageConfTh = float(sys.argv[4])
        pcdConfTh = float(sys.argv[5])
        callbackUrl = sys.argv[6]
        logger.info("data annotation container receive data:{}".format(sys.argv))
        # check param
        if not isinstance(datasetId, str):
            error_message = "datasetId {} type check failed, need str".format(datasetId)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(datasetType, str):
            error_message = "datasetType {} type check failed, need str".format(datasetType)
            logger.error(error_message)
            raise Exception(error_message)

        if datasetType not in ["1", "2", "3"]:
            error_message = "datasetType {} value check failed, need '1' for image annotation, '2' for pcd annotation, '3' for joint annotation".format(datasetType)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(datasetPath, str):
            error_message = "datasetPath {} type check failed, need str".format(datasetPath)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not datasetPath.endswith("/"):
            datasetPath = datasetPath + "/" 
        
        if not isinstance(callbackUrl, str):
            error_message = "callbackUrl {} type check failed, need str".format(callbackUrl)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(imageConfTh, float):
            error_message = "imageConfTh {} type check failed, need float".format(imageConfTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if imageConfTh < 0.000001 or imageConfTh > 0.999999: 
            error_message = "imageConfTh {} value check failed, need [0.01, 0.99]".format(imageConfTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if not isinstance(pcdConfTh, float):
            error_message = "pcdConfTh {} type check failed, need float".format(pcdConfTh)
            logger.error(error_message)
            raise Exception(error_message)
        
        if pcdConfTh < 0.00001 or pcdConfTh > 0.999999: 
            error_message = "pcdConfTh {} value check failed, need [0.01, 0.99]".format(pcdConfTh)
            logger.error(error_message)
            raise Exception(error_message)
        logger.info("data annotation container receive data checked ok")
        
        # check bucket download
        obs_client_annotation = ObsClient(access_key_id=config.annotation_ak, secret_access_key=config.annotation_sk, server=config.annotation_endpoint)
        obs_client_datacenter = ObsClient(access_key_id=config.ak, secret_access_key=config.sk, server=config.endpoint)
        logger.info("obs_client init done")
        response = obs_client_annotation.listObjects(config.annotation_bucket_name, prefix=datasetPath)
        if response.status < 300:
            logger.info("{} listObjects success".format(config.annotation_bucket_name + "/" + datasetPath))
        else:
            error_message = "{} listObjects failed".format(config.annotation_bucket_name + "/" + datasetPath)
            logger.error(error_message)
            raise Exception(error_message)
        
        #check zip file
        objects = []
        for obj in response.body.contents:
            if obj.key.endswith(".zip"):
                objects.append(obj.key)
        if len(objects) == 0:
            error_message = "no zip file in {}".format(config.bucket_name + "/" + datasetPath)
            logger.error(error_message)
            raise Exception(error_message)
        else:
            logger.info("{} zip files in {}".format(len(objects), config.bucket_name + "/" + datasetPath))

        #check upload 
        object_key = data_path_dict["pre_annotation_upload_path"] + datasetPath
        response = obs_client_annotation.putFile(config.annotation_bucket_name, object_key+"test_upload.json", file_path="./test_upload.json")
        if response.status < 300:
            logger.info("{} putFile test_upload success".format(config.annotation_bucket_name + "/" + object_key))
        else:
            error_message = "{} putFile test_upload failed".format(config.annotation_bucket_name + "/" + object_key)
            logger.error(error_message)
            raise Exception(error_message)
        
        #start annotation
        cfg_from_yaml_file("./pcd_detection/tools/cfgs/nuscenes_models/cbgs_voxel0075_voxelnext.yaml", cfg)
        demo_dataset = DemoDataset(dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False, root_path=Path("./pcd_detection/tools/000690_pandarat128_2372.pcd"), ext=".pcd", logger=logger)
        voxelnext_model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)
        voxelnext_model.load_params_from_file(filename="./pcd_detection/voxelnext_nuscenes_kernel1.pth", logger=logger, to_cpu=True)
        voxelnext_model.cuda()
        voxelnext_model.eval()
        logger.info("voxelnext model to gpu done")
        yolo_model = YOLO("./img_detection_yolo/weight/v8x_best_0717.pt")
        yolo_model.cuda()
        logger.info("yolo model to gpu done")
        annotate(voxelnext_model, yolo_model, obs_client_annotation, sys.argv, logger)
        logger.info("data annotation function done")

    except Exception as e:
        logger.error("data annotation failed because {}".format(e))
    
